import sys
import shutil
import os
import json
import subprocess
from typing import List

dlls_to_extract = {
  "gdiplus": ["gdiplus.dll"],
  "amstream": ["amstream.dll"],
  "qasf": ["qasf.dll"],
  "qedit": ["qedit.dll"],
  "qcap": ["qcap.dll"],
  "qdvd": ["qdvd.dll"],
  "quartz": ["quartz.dll"],
  "msftedit": ["msftedit.dll"],
  "mediafoundation": ["mf.dll"],
  "mfc40": ["mfc40.dll", "mfc40u.dll"]
}

# wget http://download.windowsupdate.com/msdownload/update/software/svpk/2011/02/windows6.1-kb976932-x64_74865ef2562006e51d7f9333b4a8d45b7a749dab.exe
# wget http://download.windowsupdate.com/msdownload/update/software/svpk/2011/02/windows6.1-kb976932-x86_c3516bc5c9e69fee6d9ac4f981f5b95977a8a2fa.exe
source = {
  "x64": "windows6.1-kb976932-x64_74865ef2562006e51d7f9333b4a8d45b7a749dab.exe",
  "x86": "windows6.1-kb976932-x86_c3516bc5c9e69fee6d9ac4f981f5b95977a8a2fa.exe"
}

def get_contents(arch):
  # arch: x86, x64
  filename = f"windows6.1-kb976932-{arch}.list"
  if not os.path.exists(filename):
    path = []
    for line in subprocess.check_output(["cabextract", "-l", source[arch]], text=True).splitlines():
      if line.endswith(".dll"):
        path.append(line.split("|")[-1].strip())
    
    with open(filename, "w") as f:
      f.write(json.dumps(path, indent=2))
    return path
  else:
    with open(filename, "r") as f:
      return json.load(f)

def parse_arch(path):
  if path.startswith("amd"):
    return "x64"
  else:
    return "x86"

def compress_extracted_folders():
    """
    Compress each folder in the extracted directory into separate tar.gz archives
    using maximum compression level, saving them in the compressed directory.
    """
    import os
    import tarfile

    # Create compressed directory if it doesn't exist
    if not os.path.exists('compressed'):
        os.makedirs('compressed')

    # Get all directories in extracted folder
    for item in os.listdir('extracted'):
        item_path = os.path.join('extracted', item)
        if os.path.isdir(item_path):
            # Create tar.gz archive with maximum compression
            archive_path = os.path.join('compressed', f'{item}.tar.gz')
            with tarfile.open(archive_path, 'w:gz', compresslevel=9) as tar:
                tar.add(item_path, arcname=item)

    print("All folders have been compressed and saved in the 'compressed' directory.")

def main():
    print(dlls_to_extract)

    if not shutil.which('cabextract'):
        print("Need cabextract installed")
        return

    if any([ not os.path.exists(p) for p in source.values() ]):
        print("Need windows 6.1 kb976932 download")
        return

    contents: List[str] = get_contents("x64")

    for dirname in dlls_to_extract:
        for dll in dlls_to_extract[dirname]:
            print("extracting...")
            paths_in_cab = [ p for p in contents if dll in p ]
            for path in paths_in_cab:
                dir_to_save = f"extracted/{dirname}/{parse_arch(path)}"
                os.makedirs(dir_to_save, exist_ok=True)
                subprocess.check_call( ["cabextract", "-F", path, "-d", dir_to_save, source["x64"] ])
                dll_path = f"{dir_to_save}/{path}"
                dll_name = os.path.split(dll_path)[-1]
                
                if os.path.exists(dll_path):
                    shutil.copyfile(dll_path, f"{dir_to_save}/{dll_name}")
                    shutil.rmtree(os.path.dirname(dll_path), ignore_errors=True)

    # Compress all extracted folders
    compress_extracted_folders()

if __name__ == '__main__':
    main()